#include <webots/accelerometer.h>
#include <webots/connector.h>
#include <webots/gps.h>
#include <webots/gyro.h>
#include <webots/inertial_unit.h>
#include <webots/motor.h>
#include <webots/position_sensor.h>
#include <webots/robot.h>

#include <rclcpp/rclcpp.hpp>
#include <rclcpp_action/rclcpp_action.hpp>
#include <tf2/LinearMath/Quaternion.h>
#include <webots_ros2_driver/PluginInterface.hpp>
#include <webots_ros2_driver/WebotsNode.hpp>

#include <nav_msgs/msg/odometry.hpp>
#include <sensor_msgs/msg/imu.hpp>
#include <std_msgs/msg/bool.hpp>
#include <std_msgs/msg/float64.hpp>

#include <general_interface/action/switch_connection.hpp>

namespace webots_driver
{

class WebotsDirectDriver : public webots_ros2_driver::PluginInterface
{
  private:
    using MotorTorque = std_msgs::msg::Float64;
    rclcpp::Subscription<MotorTorque>::SharedPtr leftMotorSub, rightMotorSub;
    using WheelSpeed = std_msgs::msg::Float64;
    rclcpp::Publisher<WheelSpeed>::SharedPtr leftSpeedPub, rightSpeedPub;
    rclcpp::Publisher<sensor_msgs::msg::Imu>::SharedPtr imuPub;
    rclcpp::Publisher<nav_msgs::msg::Odometry>::SharedPtr spOdomPub;
    using SwConnection = general_interface::action::SwitchConnection;
    rclcpp_action::Server<SwConnection>::SharedPtr swConnectionServer;
    std::shared_ptr<rclcpp_action::ServerGoalHandle<SwConnection>> handle;

    WbDeviceTag rightWheel;
    WbDeviceTag leftWheel;
    WbDeviceTag rightSensor;
    WbDeviceTag leftSensor;
    WbDeviceTag imu;
    WbDeviceTag gyro;
    WbDeviceTag acc;
    WbDeviceTag gps;
    WbDeviceTag recLeft;
    WbDeviceTag recRight;
    WbDeviceTag conLeft;
    WbDeviceTag conRight;
    double lastLeftPosition = 0, lastRightPosition = 0;
    double leftTorque = 0, rightTorque = 0;
    double radius = 1;
    webots_ros2_driver::WebotsNode *node;

  public:
    void init(webots_ros2_driver::WebotsNode *node,
              std::unordered_map<std::string, std::string> &parameters) override
    {
        this->node = node;
        if (!parameters.contains("radius"))
        {
            RCLCPP_ERROR(
                node->get_logger(),
                "No radius parameter specified, please specify at urdf");
            rclcpp::shutdown();
        }
        radius = std::stod(parameters.at("radius"));
        // search device
        std::string robotName = wb_robot_get_name();
        rightWheel = wb_robot_get_device("right_wheel_motor");
        leftWheel = wb_robot_get_device("left_wheel_motor");
        rightSensor = wb_motor_get_position_sensor(rightWheel);
        leftSensor = wb_motor_get_position_sensor(leftWheel);
        wb_position_sensor_enable(leftSensor, 5);
        wb_position_sensor_enable(rightSensor, 5);
        // why does this return nan
        // lastLeftPosition = wb_position_sensor_get_value(leftSensor);
        // lastRightPosition = wb_position_sensor_get_value(rightSensor);
        leftSensor = wb_motor_get_position_sensor(rightWheel);
        imu = wb_robot_get_device("imu");
        gyro = wb_robot_get_device("gyro");
        acc = wb_robot_get_device("acc");
        wb_motor_set_position(rightWheel, INFINITY);
        wb_motor_set_position(leftWheel, INFINITY);
        wb_motor_set_velocity(rightWheel, 0);
        wb_motor_set_velocity(leftWheel, 0);
        wb_gyro_enable(gyro, 5);
        wb_inertial_unit_enable(imu, 5);
        wb_accelerometer_enable(acc, 5);
        gps = wb_robot_get_device("gps");
        wb_gps_enable(gps, 5);

        conLeft = wb_robot_get_device("con_left");
        conRight = wb_robot_get_device("con_right");
        recLeft = wb_robot_get_device("rec_left");
        recRight = wb_robot_get_device("rec_right");
        wb_connector_unlock(conLeft);
        wb_connector_unlock(conRight);

        wb_connector_enable_presence(conLeft, 100);
        wb_connector_enable_presence(conRight, 100);

        if (conLeft == 0 || conRight == 0 || recLeft == 0 || recRight == 0)
        {
            RCLCPP_WARN(node->get_logger(),
                        "Connector not found. This can be ok if you are using "
                        "a robot without connectors.");
        }

        leftMotorSub = node->create_subscription<MotorTorque>(
            "left_torque", 1, [this](const MotorTorque::SharedPtr msg)
            { leftTorque = msg->data; });
        rightMotorSub = node->create_subscription<MotorTorque>(
            "right_torque", 1, [this](const MotorTorque::SharedPtr msg)
            { rightTorque = msg->data; });

        leftSpeedPub = node->create_publisher<WheelSpeed>("left_speed", 1);
        rightSpeedPub = node->create_publisher<WheelSpeed>("right_speed", 1);
        imuPub = node->create_publisher<sensor_msgs::msg::Imu>("imu", 1);
        spOdomPub =
            node->create_publisher<nav_msgs::msg::Odometry>("sp_odom", 1);

        swConnectionServer = rclcpp_action::create_server<SwConnection>(
            node, "switch_connection",
            std::bind(&WebotsDirectDriver::on_swc_goal, this,
                      std::placeholders::_1, std::placeholders::_2),
            std::bind(&WebotsDirectDriver::on_swc_cancel, this,
                      std::placeholders::_1),
            std::bind(&WebotsDirectDriver::on_swc_accept, this,
                      std::placeholders::_1));
    }
    void step() override
    {
        double dt = wb_robot_get_basic_time_step() * 1e-3;
        wb_motor_set_torque(leftWheel, leftTorque);
        wb_motor_set_torque(rightWheel, rightTorque);

        double poseLeft = wb_position_sensor_get_value(leftSensor);
        double wheelSpeedLeft = (poseLeft - lastLeftPosition) * radius / dt;
        lastLeftPosition = poseLeft;
        double poseRight = wb_position_sensor_get_value(rightSensor);
        double wheelSpeedRight = (poseRight - lastRightPosition) * radius / dt;
        lastRightPosition = poseRight;
        // RCLCPP_INFO(node->get_logger(), "Left  %f, Right  %f,dt: %f",
        //             wheelSpeedLeft, wheelSpeedRight, dt);
        WheelSpeed SpeedMsg;
        SpeedMsg.data = wheelSpeedLeft;
        leftSpeedPub->publish(SpeedMsg);
        SpeedMsg.data = wheelSpeedRight;
        rightSpeedPub->publish(SpeedMsg);

        sensor_msgs::msg::Imu imuMsg;
        const double *vals = wb_inertial_unit_get_roll_pitch_yaw(imu);
        tf2::Quaternion q;
        q.setRPY(vals[0], vals[1], vals[2]);
        imuMsg.orientation.x = q.x();
        imuMsg.orientation.y = q.y();
        imuMsg.orientation.z = q.z();
        imuMsg.orientation.w = q.w();
        imuMsg.angular_velocity.x = wb_gyro_get_values(gyro)[0];
        imuMsg.angular_velocity.y = wb_gyro_get_values(gyro)[1];
        imuMsg.angular_velocity.z = wb_gyro_get_values(gyro)[2];
        imuMsg.linear_acceleration.x = wb_accelerometer_get_values(acc)[0];
        imuMsg.linear_acceleration.y = wb_accelerometer_get_values(acc)[1];
        imuMsg.linear_acceleration.z = wb_accelerometer_get_values(acc)[2];
        imuMsg.header.stamp = rclcpp::Time(wb_robot_get_time() * 1e9);
        imuMsg.header.frame_id = "footprint";
        imuPub->publish(imuMsg);

        nav_msgs::msg::Odometry msg;
        const double *pos = wb_gps_get_values(gps);
        //msg.header.stamp = rclcpp::Time(wb_robot_get_time() * 1e9);
        msg.header.stamp = node->get_clock()->now();
        msg.header.frame_id = "map";
        msg.pose.pose.position.x = pos[0];
        msg.pose.pose.position.y = pos[1];
        msg.pose.pose.position.z = pos[2];
        msg.pose.pose.orientation.x = q.x();
        msg.pose.pose.orientation.y = q.y();
        msg.pose.pose.orientation.z = q.z();
        msg.pose.pose.orientation.w = q.w();
        msg.twist.twist.linear.x = (wheelSpeedLeft + wheelSpeedRight) / 2;
        msg.twist.twist.linear.y = 0;
        msg.twist.twist.linear.z = 0;
        msg.twist.twist.angular = imuMsg.angular_velocity;
        spOdomPub->publish(msg);

        if (conLeft != 0 && conRight != 0)
        {
            int left = wb_connector_get_presence(conLeft);
            int right = wb_connector_get_presence(conRight);
            // RCLCPP_INFO(node->get_logger(), "Connector left: %d, right:
            // %d",left, right);
            bool isConnected = left > 0 && right > 0;
            if (handle != nullptr)
            {
                if (handle->get_goal()->command ==
                        SwConnection::Goal::CONNECT &&
                    isConnected)
                {
                    auto result = std::make_shared<SwConnection::Result>();
                    result->result = SwConnection::Result::CONNECTED;
                    result->success = true;
                    if (handle->is_active())
                    {
                        handle->succeed(result);
                        handle = nullptr;
                        return;
                    }
                }
                else if (handle->get_goal()->command ==
                             SwConnection::Goal::DISCONNECT &&
                         !isConnected)
                {
                    auto result = std::make_shared<SwConnection::Result>();
                    result->result = SwConnection::Result::DISCONNECTED;
                    result->success = true;
                    if (handle->is_active())
                    {
                        handle->succeed(result);
                        handle = nullptr;
                        return;
                    }
                }

                auto fb = std::make_shared<SwConnection::Feedback>();
                if (handle->get_goal()->command == SwConnection::Goal::CONNECT)
                {
                    fb->current = SwConnection::Result::CONNECTING;
                }
                else if (handle->get_goal()->command ==
                         SwConnection::Goal::DISCONNECT)
                {
                    fb->current = SwConnection::Result::DISCONNECTING;
                }
                handle->publish_feedback(fb);
            }
        }
    }

    rclcpp_action::GoalResponse on_swc_goal(
        const rclcpp_action::GoalUUID &uuid,
        std::shared_ptr<const SwConnection::Goal> goal)
    {
        (void)uuid;
        RCLCPP_INFO(node->get_logger(), "Received goal %d", goal->command);
        return rclcpp_action::GoalResponse::ACCEPT_AND_EXECUTE;
    }

    rclcpp_action::CancelResponse on_swc_cancel(
        const std::shared_ptr<rclcpp_action::ServerGoalHandle<SwConnection>>
            goal_handle)
    {
        (void)goal_handle;
        RCLCPP_INFO(node->get_logger(), "Received request to cancel goal");
        handle = nullptr;
        return rclcpp_action::CancelResponse::ACCEPT;
    }

    void on_swc_accept(
        const std::shared_ptr<rclcpp_action::ServerGoalHandle<SwConnection>>
            goal_handle)
    {
        auto goal = goal_handle->get_goal();
        if (goal->command ==
            general_interface::action::SwitchConnection::Goal::CONNECT)
        {
            wb_connector_lock(conLeft);
            wb_connector_lock(conRight);
        }
        else if (goal->command ==
                 general_interface::action::SwitchConnection::Goal::DISCONNECT)
        {
            wb_connector_unlock(conLeft);
            wb_connector_unlock(conRight);
        }
        handle = goal_handle;
    }
};
} // namespace webots_driver
#include "pluginlib/class_list_macros.hpp"
PLUGINLIB_EXPORT_CLASS(webots_driver::WebotsDirectDriver,
                       webots_ros2_driver::PluginInterface)